package cn.you.GenghisKhan.common.utils.generator;

import cn.you.GenghisKhan.db.relation.mybatis.MybatisUtil;
import org.dom4j.Document;
import org.dom4j.io.SAXReader;
import org.mybatis.generator.api.CommentGenerator;
import org.mybatis.generator.api.IntrospectedTable;
import org.mybatis.generator.api.PluginAdapter;
import org.mybatis.generator.api.dom.java.*;
import org.mybatis.generator.api.dom.xml.Element;
import org.mybatis.generator.api.dom.xml.TextElement;
import org.mybatis.generator.api.dom.xml.XmlElement;

import java.io.File;
import java.util.HashMap;
import java.util.Iterator;
import java.util.List;
import java.util.Map;

public  class SplitTablePlugin extends PluginAdapter {
    private String tableName = "tableName";
    private Map<String,Integer> splitConfigMap=null;
    private Map<String,String> splitConfigFieldMap=null;
    private  String defaultSplitTableConfigName="/generator/splitTableConfig.xml";

    public SplitTablePlugin(){
        this.initSplitConfigXml();
    }

    /**
     * 在Exmaple类中添加tableName字段
     */
    @Override
    public boolean modelExampleClassGenerated(TopLevelClass topLevelClass, IntrospectedTable introspectedTable) {
        String tableName= MybatisUtil.getTableName(topLevelClass.getType().getShortName());

        if (null!=splitConfigMap.get(tableName)){
            Field tableNameField = new Field("tableName", PrimitiveTypeWrapper.getStringInstance());
            // 默认设置为当前的table名字
            tableNameField.setInitializationString("\""+ introspectedTable.getTableConfiguration().getTableName()+ "\"");
            tableNameField.setVisibility(JavaVisibility.PRIVATE);
            addField(topLevelClass, introspectedTable, tableNameField);
        }
        return super.modelExampleClassGenerated(topLevelClass,introspectedTable);
    }

    /**
     * 在object类中添加tableName字段
     */
    @Override
    public boolean modelBaseRecordClassGenerated(TopLevelClass topLevelClass, IntrospectedTable introspectedTable) {

        String tableName=MybatisUtil.getTableName(topLevelClass.getType().getShortName());

        if (null!=splitConfigMap.get(tableName)){
            Field tableNameField = new Field("tableName", PrimitiveTypeWrapper.getStringInstance());
            // 默认设置为当前的table名字
            tableNameField.setInitializationString("\""+ introspectedTable.getTableConfiguration().getTableName()+ "\"");
            tableNameField.setVisibility(JavaVisibility.PRIVATE);
            addField(topLevelClass, introspectedTable, tableNameField);
        }
        return super.modelBaseRecordClassGenerated(topLevelClass,introspectedTable);
    }

    /**
     * 这三个函数在分表中需要用其他函数替换，以为分表需要传入table名字，但是count函数需要处理*/
    @Override
    public boolean sqlMapDeleteByPrimaryKeyElementGenerated(XmlElement element, IntrospectedTable introspectedTable) {
        //TODO
        return true;
    };
    @Override
    public boolean sqlMapSelectByPrimaryKeyElementGenerated(XmlElement element, IntrospectedTable introspectedTable) {
        //resetSelectByPrimaryKey(element);
       // return super.sqlMapSelectByPrimaryKeyElementGenerated(element,introspectedTable);
        return true;
    };


    @Override
    public boolean sqlMapCountByExampleElementGenerated(XmlElement element, IntrospectedTable introspectedTable){
        String tableName=introspectedTable.getTableConfiguration().getTableName();
        if (null!=splitConfigMap.get(tableName)){
            resetCountByExample(element);
        }
        return super.sqlMapCountByExampleElementGenerated(element, introspectedTable);
    }

    /**
     * 在xml的SelectByExample的SQL语句添加limit
     */
    @Override
    public boolean sqlMapSelectByExampleWithoutBLOBsElementGenerated(XmlElement element, IntrospectedTable introspectedTable) {
        String tableName=introspectedTable.getTableConfiguration().getTableName();
        if (null!=splitConfigMap.get(tableName)){
            resetSelectXmlElementTableName(element);
        }
        return super.sqlMapSelectByExampleWithoutBLOBsElementGenerated(element,introspectedTable);
    }

    @Override
    public boolean sqlMapUpdateByExampleSelectiveElementGenerated(
            XmlElement element, IntrospectedTable introspectedTable) {

        String tableName=introspectedTable.getTableConfiguration().getTableName();
        if (null!=splitConfigMap.get(tableName)){
            resetUpdateXmlElementTableName(element);
        }

        return super.sqlMapUpdateByExampleWithBLOBsElementGenerated(element,
                introspectedTable);
    };
    @Override
    public boolean sqlMapUpdateByExampleWithoutBLOBsElementGenerated(
            XmlElement element, IntrospectedTable introspectedTable) {

        String tableName=introspectedTable.getTableConfiguration().getTableName();
        if (null!=splitConfigMap.get(tableName)){
            resetUpdateXmlElementTableName(element);
        }

        return super.sqlMapUpdateByExampleWithBLOBsElementGenerated(element,
                introspectedTable);
    }
    @Override
    public boolean sqlMapUpdateByPrimaryKeyWithoutBLOBsElementGenerated(
            XmlElement element, IntrospectedTable introspectedTable) {
        String tableName=introspectedTable.getTableConfiguration().getTableName();
        if (null!=splitConfigMap.get(tableName)) {
            resetUpdateXmlElementTableNameNotMapType(element);
        }
        return super.sqlMapUpdateByPrimaryKeyWithoutBLOBsElementGenerated(
                element, introspectedTable);
    };
    @Override
    public boolean sqlMapUpdateByPrimaryKeySelectiveElementGenerated(XmlElement element, IntrospectedTable introspectedTable) {
        String tableName=introspectedTable.getTableConfiguration().getTableName();
        if (null!=splitConfigMap.get(tableName)) {
            resetUpdateXmlElementTableNameNotMapType(element);
        }
        return super.sqlMapUpdateByPrimaryKeySelectiveElementGenerated(element, introspectedTable);
    };

    @Override
    public boolean sqlMapInsertElementGenerated(XmlElement element,
                                                IntrospectedTable introspectedTable) {
        String tableName=introspectedTable.getTableConfiguration().getTableName();
        if (null!=splitConfigMap.get(tableName)) {
            resetInsertXmlElementTableName(element);
        }
        return super.sqlMapInsertElementGenerated(element, introspectedTable);
    }
    @Override
    public boolean sqlMapInsertSelectiveElementGenerated(XmlElement element,
                                                         IntrospectedTable introspectedTable) {
        String tableName=introspectedTable.getTableConfiguration().getTableName();
        if (null!=splitConfigMap.get(tableName)) {
            resetInsertXmlElementTableName(element);
        }
        return super.sqlMapInsertSelectiveElementGenerated(element,
                introspectedTable);

    };

    @Override
    public boolean sqlMapDeleteByExampleElementGenerated(XmlElement element,
                                                         IntrospectedTable introspectedTable){
        String tableName=introspectedTable.getTableConfiguration().getTableName();
        if (null!=splitConfigMap.get(tableName)) {
            resetDeleteXmlElementTableName(element);
        }
        return super.sqlMapDeleteByExampleElementGenerated(element, introspectedTable);
    }



    private void resetSelectXmlElementTableName(XmlElement element) {
        List<Element> elements = element.getElements();
        Object subSentence = new TextElement("from ${" + tableName + "}");
        elements.set(3, (Element) subSentence);
    }
    private void resetInsertXmlElementTableName(XmlElement element) {

        List<Element> elements = element.getElements();
        String content = elements.get(0).getFormattedContent(0);
        String[] data = content.split(" ");
        data[2] = "${" + tableName + "}";
        TextElement subSentence = new TextElement(
                join(" ", data));
        elements.set(0, subSentence);
    }
    private void resetDeleteXmlElementTableName(XmlElement element) {

        List<Element> elements = element.getElements();
        String content = elements.get(0).getFormattedContent(0);
        String[] data = content.split(" ");
        data[2] = "${" + tableName + "}";
        TextElement subSentence = new TextElement(
                join(" ", data));
        elements.set(0, subSentence);
    }
    private void resetUpdateXmlElementTableName(XmlElement element) {
        List<Element> elements = element.getElements();
        TextElement subSentence = new TextElement("update ${record."
                + tableName + "}");
        elements.set(0, subSentence);
    }
    private void resetUpdateXmlElementTableNameNotMapType(XmlElement element) {
        List<Element> elements = element.getElements();
        TextElement subSentence = new TextElement("update ${" + tableName + "}");
        elements.set(0, subSentence);
    }

    private void resetCountByExample(XmlElement element) {
        List<Element> elements = element.getElements();
        String content = elements.get(0).getFormattedContent(0);
        String[] data = content.split(" ");
        data[3] = "${" + tableName + "}";
        TextElement subSentence = new TextElement(
                join(" ", data));
        elements.set(0, subSentence);
    }

    private void resetSelectByPrimaryKey(XmlElement element){
        List<Element> elements = element.getElements();
        Object subSentence = new TextElement("from ${" + tableName + "}");
        elements.set(2, (Element) subSentence);
    }

    public static String join(String join, String[] strAry) {
        StringBuffer sb = new StringBuffer();
        for (int i = 0; i < strAry.length; i++) {
            if (i == (strAry.length - 1)) {
                sb.append(strAry[i]);
            } else {
                sb.append(strAry[i]).append(join);
            }
        }
        return new String(sb);
    }

    /**
     * 取消验证
     */
    public boolean validate(List<String> warnings) {
        return true;
    }

    /**
     * 获取系统分隔符
     *
     * @return
     */
    protected String getSeparator() {
        return System.getProperty("line.separator");
    }
    /**
     * 添加字段，同时也添加get,set方法
     *
     * @param topLevelClass
     * @param introspectedTable
     * @param field
     */
    protected void addField(TopLevelClass topLevelClass, IntrospectedTable introspectedTable, Field field) {
        CommentGenerator commentGenerator = context.getCommentGenerator();
        // 添加Java字段
        commentGenerator.addFieldComment(field, introspectedTable);
        topLevelClass.addField(field);
        String fieldName = field.getName();
        // 生成Set方法
        Method method = new Method();
        method.setVisibility(JavaVisibility.PUBLIC);
        method.setName(generateSetMethodName(fieldName));
        method.addParameter(new Parameter(field.getType(), fieldName));
        method.addBodyLine("this." + fieldName + "=" + fieldName + ";");
        commentGenerator.addGeneralMethodComment(method, introspectedTable);
        topLevelClass.addMethod(method);

        // 生成Get方法
        method = new Method();
        method.setVisibility(JavaVisibility.PUBLIC);
        method.setReturnType(field.getType());
        method.setName(generateGetMethodName(fieldName));
        method.addBodyLine("return " + fieldName + ";");
        commentGenerator.addGeneralMethodComment(method, introspectedTable);
        topLevelClass.addMethod(method);

        //生成分表方法
        method=new Method();
        method.setVisibility(JavaVisibility.PUBLIC);
        method.setName("generatorTableName");
        method.addParameter(0,new Parameter(new FullyQualifiedJavaType("long"),"key"));
        String key=field.getInitializationString().replace("\"","");
        String line="this." + fieldName + "=" + fieldName + ";";
        if(null!=splitConfigMap.get(key)){
            line="this.tableName="+tableName+"+\"_\"+key%"+splitConfigMap.get(key)+";";
        }
        method.addBodyLine(line);
        commentGenerator.addGeneralMethodComment(method, introspectedTable);
        topLevelClass.addMethod(method);
    }

    protected static String generateGetMethodName(String fieldName) {
        return "get" + Character.toUpperCase(fieldName.charAt(0)) + fieldName.substring(1);
    }

    protected static String generateSetMethodName(String fieldName) {
        return "set" + Character.toUpperCase(fieldName.charAt(0)) + fieldName.substring(1);
    }

    /**
     * 设置分表配置文件名称
     * @param splitConfigName
     */
    protected void setSplitConfigName(String splitConfigName){
        this.defaultSplitTableConfigName=splitConfigName;
    }

    /**
     * 获得分表配置
     */
    private void   initSplitConfigXml(){
        if (null==splitConfigMap){
            splitConfigMap=new HashMap();
            splitConfigFieldMap=new HashMap();
            String path=  System.getProperty("user.dir")+"/src/main/resources/"+defaultSplitTableConfigName;
            try {
                File f = new File(path);
                if (f.exists()){
                    SAXReader reader = new SAXReader();
                    Document doc = reader.read(f);
                    org.dom4j.Element root = doc.getRootElement();
                    for(Iterator it=root.elementIterator();it.hasNext();){
                        org.dom4j.Element element = (org.dom4j.Element) it.next();
                        String name=element.attributeValue("name");
                        String size=element.attributeValue("size");
                        String field=element.attributeValue("field");
                        //System.out.println(name+" "+mod);
                        splitConfigMap.put(name,Integer.valueOf(size));
                        field= MybatisUtil.underline2Camel(field,true);
                        splitConfigFieldMap.put(name,field);
                    }
                }
            } catch (Exception e) {
                e.printStackTrace();
            }
        }
    }
}
